import numpy as np
from matplotlib import pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.decomposition import PCA
from sklearn.metrics import rand_score, adjusted_rand_score, confusion_matrix

from ml_rs.data_read import read_in


def plot_clu_data(x_pca, y_p, model):  # 聚类可视化
    centers = model.cluster_centers_
    plt.figure(figsize=(10, 8))
    class1 = np.where(y_p == 0)  # 找到y=0的位置
    class2 = np.where(y_p == 1)  # 找到y=1的位置
    class3 = np.where(y_p == 2)  # 找到y=2的位置
    class4 = np.where(y_p == 3)  # 找到y=3的位置
    class5 = np.where(y_p == 4)  # 找到y=4的位置
    p1, = plt.plot(np.ravel(x_pca[class1, 0]), np.ravel(x_pca[class1, 1]), 'or', markersize=2)
    p2, = plt.plot(np.ravel(x_pca[class2, 0]), np.ravel(x_pca[class2, 1]), 'ob', markersize=2)
    p3, = plt.plot(np.ravel(x_pca[class3, 0]), np.ravel(x_pca[class3, 1]), 'og', markersize=2)
    p4, = plt.plot(np.ravel(x_pca[class4, 0]), np.ravel(x_pca[class4, 1]), 'oc', markersize=2)
    p5, = plt.plot(np.ravel(x_pca[class5, 0]), np.ravel(x_pca[class5, 1]), 'oy', markersize=2)
    p6, = plt.plot(centers[:, 0], centers[:, 1], 'ok', markersize=4)

    plt.xlabel("X1")
    plt.ylabel("X2")
    plt.legend([p1, p2, p3, p4, p5, p6], ["y=0", "y=1", "y=2", "y=3", "y=4", "Centers"])
    plt.show()


# 降维可视化
def plot_dr_data(data1, data2):
    plt.figure(figsize=(10, 8))
    p1, = plt.plot(data1[:, 0], data1[:, 10], 'mp', markersize=8)
    p2, = plt.plot(data2[:, 0], data2[:, 10], 'ob', markersize=8)
    plt.xlabel("band1")
    plt.ylabel("band11")
    plt.legend([p1, p2], ["path", "tree"])
    plt.show()


if __name__ == '__main__':
    # 获取数据
    [dataset, label] = read_in("..\\test_data")

    # 降维
    pca = PCA(n_components=120)
    x_dr = pca.fit_transform(dataset)
    path_index = np.where(label == 'test_path')[0]
    tree_index = np.where(label == 'test_tree')[0]

    # 降维可视化
    path_feature = x_dr[path_index[0]:path_index[-1] + 1]
    tree_feature = x_dr[tree_index[0]:tree_index[-1] + 1]
    plot_dr_data(path_feature, tree_feature)

    y = np.concatenate(
        (np.full(sum(label == "test_grass"), 0),
         np.full(sum(label == "test_path"), 1),
         np.full(sum(label == "test_road"), 2),
         np.full(sum(label == "test_roof"), 3),
         np.full(sum(label == "test_tree"), 4))
    )

    # 聚类可视化
    for i in range(20):
        model = KMeans(n_clusters=5,
                       init='random',  # 随机选取初始向量
                       n_init=1).fit(x_dr)
        y_predict = model.labels_  # 输出预测标签
        plot_clu_data(x_dr, y_predict, model)
        print(confusion_matrix(y, y_predict))
        # print('RI‘s value :', rand_score(y, y_predict))
        # print('ARI‘s value :', adjusted_rand_score(y, y_predict))  #
